import os
import pickle
from contextlib import contextmanager


# 定义上下文文件路径
CONTEXT_FILE = "jupyter_like_context.pkl"

@contextmanager
def persistent_context():
    # 加载上下文
    if os.path.exists(CONTEXT_FILE):
        with open(CONTEXT_FILE, 'rb') as f:
            context = pickle.load(f)
    else:
        context = {}

    try:
        yield context  # 提供上下文给代码块
    finally:
        # 持久化上下文到硬盘
        with open(CONTEXT_FILE, 'wb') as f:
            pickle.dump(context, f)
        print("Context saved to disk.")

# 定义代码块
def cell_1(context):
    import pandas as pd
    print("Running Cell 1...")
    # context['pd'] = pd
    s = pd.Series(['a', 'b', 'c'])
    print(s)
    context['a'] = 10
    context['b'] = 20
    print(f"Set variables: a = {context['a']}, b = {context['b']}")

def cell_2(context):
    print("Running Cell 2...")
    if 'pd' in context:
        print(pd.Series(['a', 'b', 'c', 'd']))
    if 'a' in context and 'b' in context:
        context['c'] = context['a'] + context['b']
        print(f"Calculated c = a + b = {context['c']}")
    else:
        print("Error: Variables 'a' and 'b' are not defined. Run Cell 1 first.")

def cell_3(context):
    print("Running Cell 3...")
    if 'c' in context:
        context['d'] = context['c'] * 2
        print(f"Calculated d = c * 2 = {context['d']}")
    else:
        print("Error: Variable 'c' is not defined. Run Cell 2 first.")

# 主程序
if __name__ == "__main__":
    with persistent_context() as ctx:
        while True:
            print("\nAvailable Cells:")
            print("1: Cell 1 (Define variables a and b)")
            print("2: Cell 2 (Calculate c = a + b)")
            print("3: Cell 3 (Calculate d = c * 2)")
            print("q: Quit")
            
            choice = input("Enter the cell number to run (or 'q' to quit): ").strip()
            
            if choice == '1':
                cell_1(ctx)
            elif choice == '2':
                cell_2(ctx)
            elif choice == '3':
                cell_3(ctx)
            elif choice.lower() == 'q':
                print("Exiting...")
                break
            else:
                print("Invalid choice. Please try again.")